from fastapi import Request, Response
from fastapi.responses import JSONResponse
from typing import Callable, Any, Dict
import json
from starlette.middleware.base import BaseHTTPMiddleware
from app.schemas.response import success_response, error_response

class ResponseMiddleware(BaseHTTPMiddleware):
    """
    响应中间件，统一处理响应格式
    """
    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        # 调用下一个中间件或路由处理函数
        response = await call_next(request)
        
        # 如果是 JSONResponse，则尝试格式化响应
        if isinstance(response, JSONResponse):
            # 获取原始响应内容
            content = response.body.decode('utf-8')
            content_json = json.loads(content)
            
            # 检查是否已经是标准格式
            if isinstance(content_json, dict) and 'code' in content_json and 'data' in content_json and 'msg' in content_json:
                return response
            
            # 如果不是标准格式，则转换为标准格式
            status_code = response.status_code
            
            # 如果状态码不是200，则认为是错误响应
            if status_code != 200:
                error_msg = content_json.get('detail', '请求失败')
                if isinstance(error_msg, dict):
                    error_msg = json.dumps(error_msg)
                return JSONResponse(
                    content=error_response(code=status_code, msg=error_msg),
                    status_code=200  # 始终返回200，错误信息在content中
                )
            
            # 成功响应
            return JSONResponse(
                content=success_response(data=content_json),
                status_code=200
            )
        
        # 其他类型的响应直接返回
        return response 